# -*- coding: utf-8 -*-
"""
    PBAR_GenerateCargoDiagnosticImagesVer2.py
    
    Based on PBAR_GenerateCargoDiagnosticImagesVer2.py.
    Makes cargo images with markers overlayed.
    Useful for checking the locations of the marker positions in the original images.


@author: jkwong
"""

#PBAR_LiningUpRadZspecImages.py

import PBAR_Zspec, PBAR_Cargo
import numpy as np
import os, copy, glob
import matplotlib.pyplot as plt
from matplotlib import cm
from scipy import ndimage

# Set useful plot variables
plotColors = ['r', 'b', 'g', 'm', 'c', 'y', 'k'] * 10
lineStyles = ['-', '-.', ':', '_', '|'] *  10
markerTypes = ['.', 'o', 'v', '^', '<', '>', '1', '2', '3', '4', 's', 'p', '*', 'h', 'H', '+', 'x', 'D', 'd']
markerTypes = markerTypes * 2
colorList = np.array(['r', 'b', 'g', 'm', 'c', 'y'])

datasetDescription = PBAR_Zspec.ReadCargoDataDescriptionFile(r'E:\PBAR\data4\CargoSet2.txt')


dataPath = r'E:\PBAR\data4\BasicScanCargo'
plotDir = r'E:\PBAR\data4\BasicScansPlots0'


# define parameters for standard width images
acqTime = 1/60.
figureSize = (16, 10)

#################################
## DEFINE FILE NAMES

for datasetIndex, datasetDescript in enumerate(datasetDescription['scanID']):
    print(datasetIndex)
#    if datasetIndex == 1:
#        break
    if datasetIndex < 60:
        continue

    acquisitionTime = acqTime
    
    ######################
    ## READ IN DATA
    
    # Read Zspec
    filenameZspec = '%s-FDFC-All.npy' %datasetDescription['scanID'][datasetIndex]
    fullFilenameZspec = os.path.join(dataPath, filenameZspec)
    print('Loading %s' %fullFilenameZspec)
    datZspec = np.load(fullFilenameZspec)
    
    # Read Cargo
    filenameCargo = 'PBAR-%s.cargoimage' %datasetDescription['dataFile'][datasetIndex]
    fullFilenameCargo = os.path.join(dataPath, filenameCargo)
    print('Loading %s' %fullFilenameCargo)
#    datCargo= np.load(fullFilenameCargo)
    (datCargo,bpp,formatt,flag,low1,high1,low2,high2) = PBAR_Cargo.ReadCargoImage(fullFilenameCargo)


    # Read in marker files
    filenameMarker = filenameCargo.replace('cargoimage', 'cargomarker')
    fullFilenameMarker = fullFilenameCargo.replace('cargoimage', 'cargomarker')
    # some don't have marker files
    if os.path.exists(fullFilenameMarker):
#        markers = PBAR_Cargo.ReadCargoMarker(fullFilenameMarker)
        markers = PBAR_Cargo.ReadCargoMarker(fullFilenameMarker)
    else:
        markers = []

    ###################
    ## PLOTS 
    
    ## Unadjusted Images ##
    logIntensity = True
    plotMarkers = True
    
    plt.figure(figsize = figureSize)
    plt.grid()
    
    if logIntensity:
        intensity = np.log(datCargo.T)
    else:
        intensity = datCargo.T
    
    plt.imshow(intensity, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
    plt.colorbar()
    plt.title(filenameCargo + ' Original Size')
    plt.xlim((0, 2400))
    
    if plotMarkers:
        for i, mark in enumerate(markers):
            x = mark['rec_left'], mark['rec_right'], mark['rec_right'], mark['rec_left'], mark['rec_left']
            y = mark['rec_bottom'], mark['rec_bottom'], mark['rec_top'], mark['rec_top'], mark['rec_bottom']
            plt.plot(x, y, 'r')
            plt.plot(mark['x'], mark['y'], 'xr')
            if 'left' in mark:
                plt.plot(mark['left']['x'], mark['left']['y'], 'gx')
            if 'right' in mark:
                plt.plot(mark['right']['x'], mark['right']['y'], 'go')
            plt.text(max(x), min(y), mark['target'], color = 'g', fontsize = 14)
        
     # Save image to file
    plotFullFilename = os.path.join(plotDir, '%s_OriginalSize.png' %filenameCargo)
    plt.savefig(plotFullFilename)
    plt.close('all')
                
    plt.figure(figsize = figureSize)
    plt.grid()
    
    
    ###################
    ## PLOTS 
    
    ## Unadjusted Images ##
    logIntensity = False
    plotMarkers = True
    
    plt.figure(figsize = figureSize)
    plt.grid()
    
    if logIntensity:
        intensity = np.log(datCargo.T)
    else:
        intensity = datCargo.T
    
    plt.imshow(intensity, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
    plt.colorbar()
    plt.title(filenameCargo + ' Original Size')
    plt.xlim((0, 2400))
    
    if plotMarkers:
        for i, mark in enumerate(markers):
            x = mark['rec_left'], mark['rec_right'], mark['rec_right'], mark['rec_left'], mark['rec_left']
            y = mark['rec_bottom'], mark['rec_bottom'], mark['rec_top'], mark['rec_top'], mark['rec_bottom']
            plt.plot(x, y, 'r')
            plt.plot(mark['x'], mark['y'], 'xr')
            if 'left' in mark:
                plt.plot(mark['left']['x'], mark['left']['y'], 'gx')
            if 'right' in mark:
                plt.plot(mark['right']['x'], mark['right']['y'], 'go')
            plt.text(max(x), min(y), mark['target'], color = 'g', fontsize = 14)
        
     # Save image to file
    plotFullFilename = os.path.join(plotDir, '%s_OriginalSizeLOG.png' %filenameCargo)
    plt.savefig(plotFullFilename)
    plt.close('all')
                
    plt.figure(figsize = figureSize)
    plt.grid()